import numpy as np
from scipy import spatial
import matplotlib.pyplot as plt
from sko.GA import GA_TSP


def cal_time_select_cost(routine):
    routine_list = routine.tolist()
    # 计算准时的适应度
    time_select_cost_list = []
    for i in range(0, order_num):
        select_order_index = routine_list.index(i)
        temp_order = order_info[select_order_index]
        time_select_cost_list.append(temp_order)
    # print(time_select_cost_list)
    cost_time_value = []
    transform_cost_value = []
    for index, value in enumerate(time_select_cost_list):
        time_end = 0
        if index != 0:
            transform_cost_time_string = time_select_cost_list[index - 1][0] + '-' + value[0]
            temp_cost = 100000
            for i, v in enumerate(transform_cost):
                if (v['name'] == transform_cost_time_string):
                    temp_cost = v['cost']
                    break
                else:
                    temp_cost = 100000
            time_end = cost_time_value[index - 1] + value[2] + temp_cost
        else:
            temp_cost = 0
            time_end = value[2]
        cost_time_value.append(time_end)
        transform_cost_value.append(temp_cost)

    limit_time_list = []
    for index, value in enumerate(time_select_cost_list):
        temp = value.copy()
        temp.append(cost_time_value[index])
        temp.append(transform_cost_value[index])
        limit_time_list.append(temp)
    return limit_time_list


def cal_total_distance(routine):
    num_points, = routine.shape
    routine = np.concatenate([routine])
    # 计算调整工况的适应度（使用距离计算）
    distance_select_list = [distance_matrix[routine[i], routine[i + 1]] for i in range(num_points - 1)]
    # print(distance_select_list)
    distance_val = sum(distance_select_list)

    # 计算准时的适应度
    limit_time_list = cal_time_select_cost(routine)
    limit_time_value = 100
    for index, value in enumerate(limit_time_list):
        temp_value = value[3] - value[4]
        if (temp_value < 0):
            limit_time_value = 100
            break
        else:
            limit_time_value = 0
    return distance_val + limit_time_value

def create_distance_matrix(order_info, transform_cost):
    distance_matrix = []
    for index1, value1 in enumerate(order_info):
        distance_each = []
        for index2, value2 in enumerate(order_info):
            temp_key = value1[0] + '-' + value2[0]
            temp_cost = 100000
            for i, v in enumerate(transform_cost):
                if (v['name'] == temp_key):
                    temp_cost = v['cost']
                    break
                else:
                    temp_cost = 100000
            distance_each.append(temp_cost)
        distance_matrix.append(distance_each)
    return np.asarray(distance_matrix)

if __name__ == '__main__':
    order_info = [
        ['A', '001', 3, 20],
        ['A', '002', 3, 50],
        ['B', '003', 4, 50],
        ['B', '004', 4, 20],
        ['B', '005', 4, 30],
        ['A', '006', 3, 50],
        ['B', '007', 4, 50],
        ['B', '008', 4, 50],
        ['A', '009', 3, 40],
        ['B', '010', 4, 50],
    ]
    order_num = len(order_info)
    transform_cost = [
        {'name': 'A-A', 'cost': 1},
        {'name': 'B-B', 'cost': 1},
        {'name': 'A-B', 'cost': 3},
        {'name': 'B-A', 'cost': 2},
    ]
    distance_matrix = create_distance_matrix(order_info, transform_cost)

    ga_tsp = GA_TSP(func=cal_total_distance, n_dim=order_num, size_pop=100, max_iter=500, prob_mut=1)
    best_points, best_distance = ga_tsp.run()
    print(best_points, best_distance)
    best_list = cal_time_select_cost(best_points)
    print(best_list)

    fig, ax = plt.subplots(1, figsize=(10, 4))
    ax.plot(ga_tsp.generation_best_Y)
    plt.show()

    pass
